from concurrent.futures import ProcessPoolExecutor
import multiprocessing
try:
	import gdal
except:
	from osgeo import gdal
try:
	import gdalconst
except:
	from osgeo import gdalconst
try:
	import ogr
except:
	from osgeo import ogr
import numpy as np
from tqdm import tqdm
import rasterio
import datetime
import shutil
import random
import time
import math
import re
import os

# print(f"{datetime.datetime.now()}>映射网络驱动器")

# os.system(r"net use Z: \\192.168.2.172\root DELL123dell /user:root")


def progress_cb(complete, message, cb_data):
    '''Emit progress report in numbers for 10% intervals and dots for 3%'''
    if int(complete*100) % 10 == 0:
        print(f'{complete*100:.0f}', end='', flush=True)
    elif int(complete*100) % 3 == 0:
        print(f'{cb_data}', end='', flush=True)

# raster Splicing
def raster_splicing(args):
	[raster_paths, target_path] = args
	if os.path.isdir(target_path):
		folder = target_path
		name = os.path.basename(raster_paths[0]).split("_")[0]
		target_path = os.path.join(folder, f"{name}.tif")
	gdal.Warp(
		target_path, 
		raster_paths, 
		format='GTiff', 
		resampleAlg=gdalconst.GRA_Bilinear,
		multithread=10,
		callback=progress_cb,
		callback_data='.')
	print(f"\n{datetime.datetime.now()}>完成拼接：{target_path}")
	return

# raster transform vector
def raster_2_vector(args):
	[raster_path,field_name,field_value,vector_path] = args
	dataset = gdal.Open(raster_path, gdal.GA_ReadOnly)
	spatial = dataset.GetSpatialRef()
	srcband = dataset.GetRasterBand(1)
	maskband = srcband.GetMaskBand()
	vector_name = os.path.basename(vector_path)
	layer_name = os.path.splitext(vector_name)[0]
	driver = ogr.GetDriverByName('ESRI Shapefile')
	datasource = driver.CreateDataSource(vector_path)
	layer = datasource.CreateLayer(
		layer_name,srs=spatial,geom_type=ogr.wkbPolygon)
	field = ogr.FieldDefn(field_name, ogr.OFTInteger)
	layer.CreateField(field)
	gdal.Polygonize(
		srcband, 
		maskband, 
		layer, 
		field_value, 
		options=[],
		callback=progress_cb,
		callback_data='.'
		)
	datasource.Destroy()
	print(f"\n{datetime.datetime.now()}>完成矢量:",vector_path)
	return vector_path

# field filter to new shp
def field_filter_feature(args):
	[vector_path, field_name, field_value, target_path] = args
	ogr.RegisterAll()   
	datasource = ogr.Open(vector_path)
	layer = datasource.GetLayer()
	layer.SetAttributeFilter(f"{field_name}={field_value}")
	gdal.SetConfigOption("SHAPE_ENCODING", "GBK")
	driver = ogr.GetDriverByName("ESRI Shapefile")
	dest_datasource = driver.CreateDataSource(target_path)
	layer_name = os.path.splitext(os.path.basename(target_path))[0]
	dest_datasource.CopyLayer(layer, layer_name)
	datasource.Destroy()
	dest_datasource.Destroy()
	print(f"{datetime.datetime.now()}>获得矢量:",target_path)
	return

# clip raster file to tif file no same
def raster_clip(args):
	raster_path, save_folder, xcount, ycount = args
	raster_name = os.path.basename(raster_path)
	name = os.path.splitext(raster_name)[0]
	gdal.AllRegister()
	dataset = gdal.Open(raster_path)
	width,height = dataset.RasterXSize,dataset.RasterYSize
	projection = dataset.GetProjection()
	geotransform = dataset.GetGeoTransform()
	xmin, xsize = geotransform[0], geotransform[1]
	ymax, ysize = geotransform[3], geotransform[5]
	xstrid = width//xcount
	ystrid = height//ycount
	print(f"{raster_path}")
	for x in range(xcount):
		for y in range(ycount):	
			wxsize = xstrid if x!=xcount-1 else xstrid+width%xcount
			wysize = ystrid if y!=ycount-1 else ystrid+height%ycount
			left = xmin + xsize  * x * xstrid
			right = left + xsize * wxsize 
			top = ymax + ysize * y * ystrid
			bottom = top + ysize * wysize
			outputBounds = (left, bottom, right, top)
			path = os.path.join(save_folder, f"{name}_{x}_{y}.tif")
			if not os.path.exists(path):
				gdal.Warp(
					path, 
					raster_path, 
					format='GTiff', 
					outputBounds=outputBounds, 
					outputBoundsSRS=projection, 
					multithread=10)
	return

# image add
def image_add_image(args):
	base_path, path_anchors, dest_path = args
	gdal.AllRegister()
	base_dataset = gdal.Open(base_path)
	base_width = base_dataset.RasterXSize
	base_height = base_dataset.RasterYSize
	base_count = base_dataset.RasterCount
	datatype = base_dataset.GetRasterBand(1).DataType
	base_projection = base_dataset.GetProjection()
	base_geotransform = base_dataset.GetGeoTransform()
	base_bands = [
		base_dataset.GetRasterBand(index).ReadAsArray()
		for index in range(1, base_count + 1)
	] 
	for mask_path, anchor in path_anchors:
		mask_dataset = gdal.Open(mask_path)
		mask_width = mask_dataset.RasterXSize
		mask_height = mask_dataset.RasterYSize
		mask_count = mask_dataset.RasterCount
		mask_bands = [
			mask_dataset.GetRasterBand(index).ReadAsArray()
			for index in range(1, mask_count + 1)
		] 
		for array1, array2 in zip(base_bands, mask_bands):
			array = array1[
				anchor[1]:anchor[1]+mask_height,
				anchor[0]:anchor[0]+mask_width
			].copy()
			array[array2!=0] = array2[array2!=0]
			array1[
				anchor[1]:anchor[1]+mask_height,
				anchor[0]:anchor[0]+mask_width
			] = array
		mask_dataset.FlushCache()
	driver = gdal.GetDriverByName('GTiff')
	dest_dataset = driver.Create(
		dest_path,
		base_width,
		base_height,
		base_count, 
		datatype)
	dest_dataset.SetProjection(base_projection)
	dest_dataset.SetGeoTransform(base_geotransform)
	for index in range(1, base_count + 1):
		dest_band = dest_dataset.GetRasterBand(index)
		array = base_bands[index-1]
		dest_band.WriteArray(array)
	base_dataset.FlushCache()
	dest_dataset.FlushCache()
	return

# label add
def label_add_image(args):
	base_path, path_anchors, dest_path = args

	gdal.AllRegister()
	base_dataset = gdal.Open(base_path)
	base_width = base_dataset.RasterXSize
	base_height = base_dataset.RasterYSize
	base_count = base_dataset.RasterCount
	datatype = base_dataset.GetRasterBand(1).DataType
	base_projection = base_dataset.GetProjection()
	base_geotransform = base_dataset.GetGeoTransform()
	base_bands = [
		base_dataset.GetRasterBand(index).ReadAsArray()
		for index in range(1, base_count + 1)
	] 
	for mask_path, anchor in path_anchors:
		mask_dataset = gdal.Open(mask_path)
		mask_width = mask_dataset.RasterXSize
		mask_height = mask_dataset.RasterYSize
		mask_count = mask_dataset.RasterCount
		mask_bands = [
			mask_dataset.GetRasterBand(index).ReadAsArray()
			for index in range(1, mask_count + 1)
		] 
		for array1, array2 in zip(base_bands, mask_bands):
			array = array1[
				anchor[1]:anchor[1]+mask_height,
				anchor[0]:anchor[0]+mask_width
			].copy()
			array[array2==0] = array2[array2==0]
			array1[
				anchor[1]:anchor[1]+mask_height,
				anchor[0]:anchor[0]+mask_width
			] = array
		mask_dataset.FlushCache()

	driver = gdal.GetDriverByName('GTiff')
	dest_dataset = driver.Create(
		dest_path,
		base_width,
		base_height,
		base_count, 
		datatype)
	dest_dataset.SetProjection(base_projection)
	dest_dataset.SetGeoTransform(base_geotransform)
	for index in range(1, base_count + 1):
		dest_band = dest_dataset.GetRasterBand(index)
		array = base_bands[index-1]
		dest_band.WriteArray(array)

	base_dataset.FlushCache()
	dest_dataset.FlushCache()
	return

# get raster mask raster file
def raster_2_mask(args):
	base_path, dest_path, pixel_value = args
	gdal.AllRegister()
	base_dataset = gdal.Open(base_path)
	width = base_dataset.RasterXSize
	height = base_dataset.RasterYSize
	base_projection = base_dataset.GetProjection()
	base_geotransform = base_dataset.GetGeoTransform()
	datatype = base_dataset.GetRasterBand(1).DataType
	array = np.ones([height,width]) * pixel_value
	driver = gdal.GetDriverByName('GTiff')
	dest_dataset = driver.Create(
		dest_path,width,height,1,datatype)
	dest_dataset.SetProjection(base_projection)
	dest_dataset.SetGeoTransform(base_geotransform)
	dest_band = dest_dataset.GetRasterBand(1)
	dest_band.WriteArray(array)
	base_dataset.FlushCache()
	dest_dataset.FlushCache()
	return

# get positive mask raster file 
def get_positive_mask(args):
	base_path, dest_path = args
	gdal.AllRegister()
	base_dataset = gdal.Open(base_path)
	width = base_dataset.RasterXSize
	height = base_dataset.RasterYSize
	base_count = base_dataset.RasterCount
	base_projection = base_dataset.GetProjection()
	base_geotransform = base_dataset.GetGeoTransform()
	datatype = base_dataset.GetRasterBand(1).DataType
	for index in range(1,base_count+1):
		band = base_dataset.GetRasterBand(index)
		if index ==1:
			array = band.ReadAsArray()
			array[array!=0] = 1
		else:
			temp = band.ReadAsArray()
			temp[temp!=0] = 1
			array = array+temp
			array[array!=0] = 1
	array[array==0] = 254
	array[array==1] = 0
	array[array==254] = 1
	driver = gdal.GetDriverByName('GTiff')
	dest_dataset = driver.Create(
		dest_path,width,height,1,datatype)
	dest_dataset.SetProjection(base_projection)
	dest_dataset.SetGeoTransform(base_geotransform)
	dest_band = dest_dataset.GetRasterBand(1)
	dest_band.WriteArray(array)
	base_dataset.FlushCache()
	dest_dataset.FlushCache()
	return

# tif horizontal flip
def horizontal_flip_tif(args):
	base_path, dest_path = args
	gdal.AllRegister()
	base_dataset = gdal.Open(base_path)
	width = base_dataset.RasterXSize
	height = base_dataset.RasterYSize
	base_count = base_dataset.RasterCount
	base_projection = base_dataset.GetProjection()
	base_geotransform = base_dataset.GetGeoTransform()
	datatype = base_dataset.GetRasterBand(1).DataType
	arrays = []
	for index in range(1,base_count+1):
		band = base_dataset.GetRasterBand(index)
		array = band.ReadAsArray()
		arrays.append(np.fliplr(array))
	driver = gdal.GetDriverByName('GTiff')
	dest_dataset = driver.Create(
		dest_path,width,height,base_count,datatype)
	dest_dataset.SetProjection(base_projection)
	dest_dataset.SetGeoTransform(base_geotransform)
	for index in range(1,base_count+1):
		band = dest_dataset.GetRasterBand(index)
		array = arrays[index-1]
		band.WriteArray(array)
	base_dataset.FlushCache()
	dest_dataset.FlushCache()
	return

# clip raster by windows
def window_clip_raster(args):
	raster_path, save_folder, window, step = args

	raster_name = os.path.basename(raster_path)
	name = os.path.splitext(raster_name)[0]

	gdal.AllRegister()
	dataset = gdal.Open(raster_path)
	width,height = dataset.RasterXSize,dataset.RasterYSize
	projection = dataset.GetProjection()
	geotransform = dataset.GetGeoTransform()
	xmin, xsize = geotransform[0], geotransform[1]
	ymax, ysize = geotransform[3], geotransform[5]

	xcount = math.ceil((width-window[0])/step[0])
	ycount = math.ceil((height-window[1])/step[1])

	for x in range(xcount):
		for y in range(ycount):	
			left = xmin + x * step[0] * xsize
			right = left + window[0] * xsize
			top = ymax + y * step[1] * ysize
			bottom = top + window[1] * ysize
			outputBounds = (left, bottom, right, top)
			path = os.path.join(save_folder, f"{name}_{x}_{y}.tif")
			if not os.path.exists(path):
				gdal.Warp(
					path, 
					raster_path, 
					format='GTiff', 
					outputBounds=outputBounds, 
					outputBoundsSRS=projection, 
					multithread=10)
	return

# raster transform png
def raster_transform_png(args):
	raster_path, png_path = args
	print(f"{png_path}")
	if not os.path.exists(png_path):
		gdal.Translate(png_path,raster_path,format='PNG')
	return


def delete_area_geometry(vector_path):
	pixel_size = (5.36441802978515625e-06)*(5.36441802978515625e-06)
	threshold = 80000 * pixel_size
	datasource = ogr.Open(vector_path,2)
	layer = datasource.GetLayer()
	fids = []
	feature = layer.GetNextFeature()
	while feature:
		geometry = feature.geometry()
		area = geometry.GetArea()
		fid = feature.GetFID()
		if area <= threshold:
			fids.append(fid)
		feature = layer.GetNextFeature()
	layer.ResetReading()
	for fid in tqdm(fids,desc=f'正在删除({vector_path})小面积要素：'):
		layer.DeleteFeature(fid)	
	datasource.Destroy()
	return


def delete_area_geometry_batch(shp_dir):
	args_list = []
	for name in os.listdir(shp_dir):
		if re.match("^.*\.shp$", name):
			vector_path = os.path.join(shp_dir, name)
			args_list.append(vector_path)
	pool = multiprocessing.Pool()
	pool.map(delete_area_geometry, args_list)
	pool.close()
	pool.join()
	return


def merge_vector(args):
	shp_paths, vector_path = args
	# vector_name = os.path.basename(vector_path)
	# layer_name = os.path.splitext(vector_name)[0]
	if len(shp_paths) == 0:
		print(vector_path, "fial !")
		return
	gdal.SetConfigOption("SHAPE_ENCODING", "GBK")
	driver = ogr.GetDriverByName('ESRI Shapefile')
	driver.CopyDataSource(ogr.Open(shp_paths.pop(0)),vector_path)

	features = []
	for shp_path in shp_paths:
		datasource = ogr.Open(shp_path)
		layer = datasource.GetLayer()
		feature = layer.GetNextFeature()
		while feature:
			features.append(feature)
			feature = layer.GetNextFeature()
		datasource.Destroy()

	datasource = ogr.Open(vector_path,2)
	layer = datasource.GetLayer()
	for feature in tqdm(features,desc=f"正在合并矢量{vector_path}:"):
		layer.CreateFeature(feature)
	datasource.SyncToDisk()
	datasource.Destroy()
	return


def merge_vector_bacth():
	
	source_dir = r"Z:\orange\20220623\第二轮解译结果"
	result_shp = r"Z:\orange\20220623\result.shp"

	shp_paths = []
	for file_name in os.listdir(source_dir):
		if re.match(".*\.shp", file_name):
			shp_path = os.path.join(source_dir, file_name)
			shp_paths.append(shp_path)

	merge_vector([shp_paths, result_shp])
	# pool = multiprocessing.Pool()
	# pool.map(merge_vector, args_list)
	# pool.close()
	# pool.join()
	return


def raster_clip():
	out_raster=r"F:\新建文件夹(2)\0020_0014_TimeOne_test.tif"
	in_raster=r"F:\新建文件夹(2)\0020_0014_TimeOne.tif"
	top,bottom,left,right=[0,0,10,10]

	dataset = gdal.Open(in_raster)
	width = dataset.RasterXSize
	height = dataset.RasterYSize
	geotransform = dataset.GetGeoTransform()
	minX = geotransform[0] + (-left) * geotransform[1]
	maxY = geotransform[3] + (-top) * geotransform[5]
	maxX = geotransform[0] + (width+right) * geotransform[1]
	minY = geotransform[3] + (height+bottom) * geotransform[5]

	outputBounds = (minX, minY, maxX, maxY)
	print(outputBounds)
	gdal.Warp(
		out_raster,
        in_raster,
        format='GTiff',
		outputBounds=outputBounds,
		cropToCutline=True,  # 按掩膜图层范围裁剪
		# dstNodata=-9999,
	
		)
	return


def tesk():
	tif=r'E:\新建文件夹(2)\0020_0014_TimeOne.tif'

	# 根据格网编码获取对应的
	SampleFileSrc = rasterio.open(tif)

	profile = SampleFileSrc.profile
	Width = SampleFileSrc.width
	Height = SampleFileSrc.height

	SampleFileRaster = SampleFileSrc.read()  # 读取所有波段

	SampleFileRaster = SampleFileRaster.astype('uint16')

	CurBandData2 = SampleFileRaster[0].reshape(Width * Height)

	list3 = CurBandData2.reshape(Width, Height)
	list3 = np.array([list3])

	cur_savedir='E:\新建文件夹(2)'

	OutTif = os.path.join(cur_savedir,  'Predict0805' + '.tif')
	profile.update(dtype=rasterio.uint16, count=1, nodata=None)
	with rasterio.open(OutTif, mode='w', **profile) as dst:
		dst.write(list3)
	return


# start
def main():
	''''''
	shp_dir = r"G:\shp"
	delete_area_geometry_batch(shp_dir)
	'''raster 2 png'''
	
	# base_folder = r"Z:\orange\clip_raster\G50F013017"
	# dest_folder = r"Z:\orange\clip_png\G50F013017"
	# args_list = []
	# for name in os.listdir(base_folder):
	# 	new_name = name.split(".")[0]+".png"
	# 	args_list.append([
	# 		os.path.join(base_folder, name),
	# 		os.path.join(dest_folder, new_name)
	# 	])

	# source_folder = r"Z:\orange\clip_raster"
	# target_folder = r"Z:\orange\clip_png"
	# args_list = []
	# for folder_name in os.listdir(source_folder):
	# 	image_folder = os.path.join(target_folder, folder_name)
	# 	if not os.path.exists(image_folder):
	# 		os.makedirs(image_folder)
	# 	tif_folder = os.path.join(source_folder, folder_name)
	# 	for file_name in os.listdir(tif_folder):
	# 		png_name = os.path.splitext(file_name)[0]+".png"
	# 		png_path = os.path.join(image_folder, png_name)
	# 		tif_path = os.path.join(tif_folder, file_name)
	# 		args_list.append([tif_path, png_path])

	# pool = multiprocessing.Pool()
	# pool.map(raster_transform_png, args_list)
	# pool.close()
	# pool.join()
	'''clip raster by windows'''
	# source_dir = r"Z:\orange\xiugai_shp\dataset_image"
	# target_dir = r"Z:\orange\xiugai_shp\dataset\image"
	# for name in os.listdir(source_dir):
	# 	window_clip_raster([
	# 		os.path.join(source_dir, name),
	# 		target_dir,
	# 		[512, 512],
	# 		[512, 512]
	# 	])
	'''flip tif'''
	# source_folder = r"Z:\orange\xiugai_shp\positive_tif"
	# target_folder = r"Z:\orange\xiugai_shp\flip_horizontal"
	# for name in os.listdir(source_folder):
	# 	new_name = name.replace(".tif", "_1.tif")
	# 	args = [
	# 		os.path.join(source_folder, name),
	# 		os.path.join(target_folder, new_name)
	# 	]
	# 	horizontal_flip_tif(args)
	'''raster add and mask add'''	
	# source_image_folder = r"Z:\orange\xiugai_shp\source_clip"
	# source_label_folder = r"Z:\orange\xiugai_shp\source_mask"
	# positive_image_folder = r"Z:\orange\xiugai_shp\new_positive"
	# positive_label_folder = r"Z:\orange\xiugai_shp\mask_transform"
	# save_image_folder = r"Z:\orange\xiugai_shp\dataset_image"
	# save_lable_folder = r"Z:\orange\xiugai_shp\dataset_label"
	# for name_1 in os.listdir(source_image_folder):
	# 	path_anchors_1 =[]
	# 	path_anchors_2 =[]
	# 	for name_2 in os.listdir(positive_image_folder):
	# 		anchor = (
	# 			random.randint(0,7500-500),
	# 			random.randint(0,5000-500))
	# 		path_anchors_1.append([
	# 			os.path.join(positive_image_folder,name_2),anchor])
	# 		path_anchors_2.append([
	# 			os.path.join(positive_label_folder,name_2),anchor])
	# 	image_add_image([
	# 		os.path.join(source_image_folder, name_1),
	# 		path_anchors_1,
	# 		os.path.join(save_image_folder, name_1)])
	# 	label_add_image([
	# 		os.path.join(source_label_folder, name_1),
	# 		path_anchors_2,
	# 		os.path.join(save_lable_folder, name_1)])
	'''get image mask'''
	# source_folder = r"Z:\orange\xiugai_shp\source_clip"
	# target_folder = r"Z:\orange\xiugai_shp\source_mask"
	# args_list = []
	# for name in os.listdir(source_folder):
	# 	args = [
	# 		os.path.join(source_folder, name), 
	# 		os.path.join(target_folder, name),
	# 		1
	# 	]
	# 	args_list.append(args)
	# with ProcessPoolExecutor(4) as executor:
	# 	executor.map(raster_2_mask, args_list)
	'''filter feature by field'''
	# source_dir = r"Z:\orange\20220412-SHP"
	# target_dir = r"Z:\orange\20220412-R-SHP"
	# for name in os.listdir(source_dir):
	# 	if re.match(".*\.shp$", name):
	# 		vector_path = os.path.join(source_dir, name)
	# 		target_path = os.path.join(target_dir, name)
	# 		field_filter_feature([
	# 			vector_path, "ID", 0, target_path])

	'''raster transform vector'''
	# source_dir = r"Z:\orange\20220412-TIF"
	# target_dir = r"Z:\orange\20220412-SHP"
	# args_list = []
	# field_name, field_value = 'ID', 0
	# for name in os.listdir(source_dir):
	# 	if re.match("^.*\.png$", name):
	# 		path = os.path.join(source_dir, name)
	# 		args = [path,target_dir,field_name,field_value]
	# 		args_list.append(args)

	# pool = multiprocessing.Pool(5)
	# pool.map(raster_2_vector, args_list)
	# pool.close()
	# pool.join()
	'''raster transform vector and filter feature by field'''
	# source_dir = r"Z:\orange\result\pseudo_color_prediction"
	# temp_dir = r"Z:\orange\20220412-SHP"
	# target_dir = r"Z:\orange\20220412-R-SHP"
	# field_name, field_value = 'ID', 0
	
	# for name in os.listdir(source_dir):
	# 	if re.match("^.*\.png$", name):
			
	# 		raster_path = os.path.join(source_dir, name)
	# 		vector_name = os.path.splitext(name)[0]+".shp"

	# 		temp_path = os.path.join(temp_dir,vector_name)
	# 		# if not os.path.exists(temp_path):
	# 		args = [raster_path,field_name,field_value,temp_path]
	# 		raster_2_vector(args)
			
			
	# 		target_path=os.path.join(target_dir,vector_name)
	# 		# if not os.path.exists(target_path):
	# 		args=[temp_path,field_name,field_value,target_path]
	# 		field_filter_feature(args)
			
	'''raster windows clip, last one big'''
	# base_folder = r"E:\新建文件夹"
	# save_folder = r"E:\Image2"
	# args_list = []
	# for name in os.listdir(base_folder):
	# 	path = os.path.join(base_folder, name)
	# 	args = [path,save_folder,3,3]
	# 	args_list.append(args)

	# pool = multiprocessing.Pool()
	# pool.map(raster_clip, args_list)
	# pool.close()
	# pool.join()
	'''raster Splicing'''
	# source_dir = r"Z:\orange\xiugai_shp"
	# target_dir = r"Z:\orange\fengting_class"
	# for folder_name in os.listdir(source_dir):
	# 	folder = os.path.join(source_dir, folder_name)
	# 	paths = []
	# 	for name in os.listdir(folder):
	# 		if re.match(".*\.png$", name):
	# 			paths.append(os.path.join(folder, name))
	# 	target_path = os.path.join(target_dir, f"{folder_name}.tif")
	# 	raster_splicing([paths, target_path])
	return


if __name__ == "__main__":
	main()
	# tesk()
	# raster_clip()
	# merge_vector_bacth()